#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created by PyCharm.

@Date    : Sat Feb 06 2021 
@Time    : 11:57:03
@File    : repvgg.py
@Author  : alpha
"""

import torch
import torch.nn as nn

from src.repconv import RepConv2D


class RepVGG18(nn.Module):

    def __init__(self, num_classes):
        super(RepVGG18, self).__init__()
        self.conv1 = RepConv2D(3, 32, stride=2)


        self.stage2_1 = RepConv2D(32, 64, stride=2)
        self.stage2_2 = RepConv2D(64, 64, stride=1)
        self.stage2_3 = RepConv2D(64, 64, stride=1)
        self.stage2_4 = RepConv2D(64, 64, stride=1)

        self.stage3_1 = RepConv2D(64, 128, stride=2)
        self.stage3_2 = RepConv2D(128, 128, stride=1)
        self.stage3_3 = RepConv2D(128, 128, stride=1)
        self.stage3_4 = RepConv2D(128, 128, stride=1)

        self.stage4_1 = RepConv2D(128, 256, stride=2)
        self.stage4_2 = RepConv2D(256, 256, stride=1)
        self.stage4_3 = RepConv2D(256, 256, stride=1)
        self.stage4_4 = RepConv2D(256, 256, stride=1)

        self.stage5_1 = RepConv2D(256, 512, stride=2)
        self.stage5_2 = RepConv2D(512, 512, stride=1)
        self.stage5_3 = RepConv2D(512, 512, stride=1)
        self.stage5_4 = RepConv2D(512, 512, stride=1)

        self.avepool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)


    def forward(self, x):
        conv1 = self.conv1(x)

        stage2 = self.stage2_4(self.stage2_3(self.stage2_2(self.stage2_1(conv1))))
        stage3 = self.stage3_4(self.stage3_3(self.stage3_2(self.stage3_1(stage2))))
        stage4 = self.stage4_4(self.stage4_3(self.stage4_2(self.stage4_1(stage3))))
        stage5 = self.stage5_4(self.stage5_3(self.stage5_2(self.stage5_1(stage4))))

        avepool = self.avepool(stage5)[:, :, 0, 0]
        fc = self.fc(avepool)
        return fc